package com.wfh.aiagent.rag;

import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;
import org.springframework.util.StreamUtils;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;

/**
 * @author fenghuanwang
 */
@Component
public class SemanticChunker {

    private final EmbeddingModel embeddingModel;
    private final double similarityThreshold = 0.75;
    private final int windowSize = 300;
    private final int windowOverlap = 20;

    public SemanticChunker(@Qualifier("dashscopeEmbeddingModel") EmbeddingModel embeddingModel) {
        this.embeddingModel = embeddingModel;
    }

    public  List<Document> split(Resource resource) {
        if (!resource.exists()) {
            return Collections.emptyList();
        }

        String text;
        try {
            text = StreamUtils.copyToString(resource.getInputStream(), StandardCharsets.UTF_8);
        } catch (IOException e) {
            throw new RuntimeException("Failed to read resource: " + resource, e);
        }

        if (text == null || text.trim().isEmpty()) {
            return Collections.emptyList();
        }

        return semanticSplit(text, resource);
    }

    /**
     * 语义化分割
     * @param text
     * @param resource
     * @return
     */
    private List<Document> semanticSplit(String text, Resource resource) {
        List<String> windows = slidingWindow(text, windowSize, windowOverlap);
        if (windows.isEmpty()) {
            return Collections.emptyList();
        }

        final int MAX_BATCH_SIZE = 25;
        List<float[]> allEmbeddings = new ArrayList<>(windows.size());

        // 🔥 关键：分批调用嵌入模型，每批最多 25 个
        for (int start = 0; start < windows.size(); start += MAX_BATCH_SIZE) {
            int end = Math.min(start + MAX_BATCH_SIZE, windows.size());
            List<String> batch = windows.subList(start, end);

            // 调用 DashScope 嵌入（单次 ≤25）
            EmbeddingResponse response = embeddingModel.embedForResponse(batch);
            for (var result : response.getResults()) {
                allEmbeddings.add(result.getOutput()); // float[]
            }
        }

        // 确保嵌入数量与窗口数量一致
        if (allEmbeddings.size() != windows.size()) {
            throw new IllegalStateException("Embedding count mismatch: windows=" + windows.size() + ", embeddings=" + allEmbeddings.size());
        }

        // 聚合语义块
        List<List<String>> chunks = new ArrayList<>();
        List<String> currentChunk = new ArrayList<>();
        currentChunk.add(windows.get(0));

        for (int i = 1; i < windows.size(); i++) {
            double similarity = cosineSimilarity(allEmbeddings.get(i - 1), allEmbeddings.get(i));
            if (similarity >= similarityThreshold) {
                currentChunk.add(windows.get(i));
            } else {
                chunks.add(new ArrayList<>(currentChunk));
                currentChunk.clear();
                currentChunk.add(windows.get(i));
            }
        }
        if (!currentChunk.isEmpty()) {
            chunks.add(currentChunk);
        }

        // 构建 Document
        List<Document> documents = new ArrayList<>();
        String filename = Optional.ofNullable(resource.getFilename()).orElse("unknown");
        String uri = "";
        try {
            uri = resource.getURI().toString();
        } catch (Exception ignored) {
        }

        for (List<String> chunk : chunks) {
            String content = String.join("", chunk);
            Map<String, Object> metadata = Map.of("source", filename, "uri", uri);
            documents.add(new Document(content, metadata));
        }

        return documents;
    }

    /**
     * 滑动窗口
     * @param text
     * @param size
     * @param overlap
     * @return
     */
    private static List<String> slidingWindow(String text, int size, int overlap) {
        List<String> result = new ArrayList<>();
        int step = Math.max(1, size - overlap);
        for (int i = 0; i < text.length(); i += step) {
            int end = Math.min(i + size, text.length());
            result.add(text.substring(i, end));
        }
        return result;
    }

    /**
     * 计算余弦相似度
     * @param a
     * @param b
     * @return
     */
    private double cosineSimilarity(float[] a, float[] b) {
        double dot = 0.0, normA = 0.0, normB = 0.0;
        for (int i = 0; i < a.length; i++) {
            dot += a[i] * b[i];
            normA += a[i] * a[i];
            normB += b[i] * b[i];
        }
        if (normA == 0 || normB == 0){
            return 0.0;
        }
        return dot / (Math.sqrt(normA) * Math.sqrt(normB));
    }
}